#include "fusecache.h"

#include "host_access.h"
#include "minip.h"

#include "kcache.h"
#include <wdm.h>

typedef struct FUSE_CACHE {
    KstlCache*   Handle;
    FUSE_CONTEXT Context;
    FAST_MUTEX   Lock;
} FUSE_CACHE, *PFUSE_CACHE, **PPFUSE_CACHE;

static void
ForgetRequest(void* Context, void* Item)
{
    PFUSE_CONTEXT Ctx        = Context;
    if (!HostShareIsConnected(Ctx->pShare)) {
        return;
    }

    UINT64 FuseNodeId = *(PUINT64)Item;
    FuseForgetRequest(Ctx->pDevExt, Ctx->pShare->ShareId, FuseNodeId, 1);
    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCache forgetRequest, value: %zu\n", FuseNodeId));
}

FUSE_CACHE_RESULT
FuseCacheInit(PPFUSE_CACHE ppCache, FUSE_CONTEXT context, size_t maxElements, unsigned long long ttlNs)
{
    if (!ppCache) {
        return INVALID_HANDLE;
    }

    if (*ppCache) {
        HWFS_DBG_PRINT(DBG_WARN, DBG_FUSE_CACHE, ("FuseCacheInit: cache %p already initialized\n", *ppCache));
        return SECOND_INITIALIZATION;
    }

    PFUSE_CACHE pCache = ExAllocatePool2(POOL_FLAG_NON_PAGED, sizeof(FUSE_CACHE), 'FCI');
    if (!pCache) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_FUSE_CACHE, ("FuseCacheInit: cache alloc failed\n"));
        return ALLOCATION_ERROR;
    }

    pCache->Context               = context;
    KstlDeallocator   Deallocator = {.deallocatorContext = &pCache->Context, .deallocate = ForgetRequest};
    KstlCacheSettings Settings    = {.itemSize = sizeof(UINT64),
                                        .maxElements = maxElements, .ttl = ttlNs, .dealloc = Deallocator};
    
    pCache->Handle = KstlCacheInit(Settings);
    if (!pCache->Handle) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_FUSE_CACHE, ("FuseCacheInit: KstlCacheInit() failed\n", *ppCache));
        FuseCacheRelease(pCache);
        return ALLOCATION_ERROR;
    }

    ExInitializeFastMutex(&pCache->Lock);

    *ppCache = pCache;

    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheInit %p\n", pCache));

    return SUCCESS;
}

void
FuseCacheRelease(PFUSE_CACHE pCache)
{
    if (!pCache) {
        return;
    }

    if (pCache->Handle) {
        KstlCacheRelease(pCache->Handle);
    }

    ExFreePool(pCache);

    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheRelease %p\n", pCache));
}

FUSE_CACHE_RESULT
FuseCacheAdd(PFUSE_CACHE pCache, PCHAR Key, UINT64 Item)
{
    if (!pCache) {
        return INVALID_HANDLE;
    }

    ExAcquireFastMutex(&pCache->Lock);

    int result = KstlCacheAdd(pCache->Handle, Key, strlen(Key), &Item);
    if (result) {
        HWFS_DBG_PRINT(DBG_ERR, DBG_FUSE_CACHE, ("FuseCacheAdd failed for key: %s, value: %zu\n", Key, Item));
        ExReleaseFastMutex(&pCache->Lock);
        return FAIL;
    }

    HWFS_DBG_PRINT(DBG_TRACE, DBG_FUSE_CACHE, ("FuseCacheAdd key: %s, value: %zu\n", Key, Item));
    ExReleaseFastMutex(&pCache->Lock);
    return SUCCESS;
}

FUSE_CACHE_RESULT
FuseCacheGet(PFUSE_CACHE pCache, PCHAR Key, PUINT64 Item)
{
    if (!pCache) {
        return INVALID_HANDLE;
    }

    ExAcquireFastMutex(&pCache->Lock);

    int result = KstlCacheGet(pCache->Handle, Key, strlen(Key), Item);
    if (result) {
        HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheGet no such key: %s\n", Key));
        ExReleaseFastMutex(&pCache->Lock);
        return NO_KEY;
    }

    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheGet existing key: %s, value: %zu\n", Key, *Item));
    ExReleaseFastMutex(&pCache->Lock);
    return SUCCESS;
}

FUSE_CACHE_RESULT
FuseCacheRemove(PFUSE_CACHE pCache, PCHAR Key)
{
    if (!pCache) {
        return INVALID_HANDLE;
    }

    ExAcquireFastMutex(&pCache->Lock);

    int result = KstlCacheRemove(pCache->Handle, Key, strlen(Key));
    if (result) {
        ExReleaseFastMutex(&pCache->Lock);
        return NO_KEY;
    }

    ExReleaseFastMutex(&pCache->Lock);
    return SUCCESS;
}


FUSE_CACHE_RESULT
FuseCacheClear(PFUSE_CACHE pCache)
{
    if (!pCache) {
        return INVALID_HANDLE;
    }

    ExAcquireFastMutex(&pCache->Lock);

    KstlCacheClear(pCache->Handle);

    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheClear\n"));
    ExReleaseFastMutex(&pCache->Lock);
    return SUCCESS;
}

FUSE_CACHE_RESULT
FuseCacheRemoveAllWithPrefix(PFUSE_CACHE pCache, PCHAR PathPrefix)
{
    if (!pCache) {
        return INVALID_HANDLE;
    }

    ExAcquireFastMutex(&pCache->Lock);

    size_t prefixSize = strlen(PathPrefix);
    if (PathPrefix[prefixSize - 1] == '\\') {
        prefixSize -= 1;
    }

    KC_FOR_EACH(pCache->Handle, it)
    {
        size_t keySize = KstlCacheIteratorKeySize(it);
        if (keySize < prefixSize) {
            continue;
        }

        PCHAR key = KstlCacheIteratorKey(it);
        if (strncmp(PathPrefix, key, prefixSize) != 0) {
            continue;
        }

        if (keySize > prefixSize && key[prefixSize + 1] != '\\') {
            continue;
        }

        KstlCacheRemoveIt(pCache->Handle, it);
    }
    HWFS_DBG_PRINT(DBG_INFO, DBG_FUSE_CACHE, ("FuseCacheRemoveAllWithPrefix\n"));

    ExReleaseFastMutex(&pCache->Lock);
    return SUCCESS;
}